#!/usr/bin/env python

# Copyright (c) 2018 Intel Labs.
# authors: German Ros (german.ros@intel.com)
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.

""" This module contains a local planner to perform
low-level waypoint following based on PID controllers. """

from collections import deque
from enum import Enum
import math
import numpy as np

import carla

from agents.navigation.controller import VehiclePIDController
from agents.navigation.stanley_controller import StanleyController

from agents.tools.misc import compute_magnitude_angle, compute_distance, clip_angle
from agents.tools.misc import distance_vehicle, draw_waypoints, draw_way, draw_speed_with_tf, draw_transforms
from agents.tools.misc import get_speed, get_acceleration, get_nearest_waypoint_index, get_curvature, get_nearest_location_index
from agents.tools.vehicle import MotionStatus
from agents.tools.bezier_curv import GetWaveBezier

class RoadOption(Enum):
    """
    RoadOption represents the possible topological configurations
    when moving from a segment of lane to other.
    """
    VOID = -1
    LEFT = 1
    RIGHT = 2
    STRAIGHT = 3
    LANEFOLLOW = 4
    CHANGELANELEFT = 5
    CHANGELANERIGHT = 6
    CROSSWALK = 7

class RTKPlanner(object):
    """
    RTKPlanner implements the basic behavior of following a trajectory
    of waypoints that is generated on-the-fly.
    The low-level motion of the vehicle is computed by using two PID controllers,
    one is used for the lateral control
    and the other for the longitudinal control (cruise speed).

    When multiple paths are available (intersections)
    this local planner makes a random choice.
    """

    # Minimum distance to target waypoint as a percentage
    # (e.g. within 80% of total distance)

    # FPS used for dt
    FPS = 20

    def __init__(self, agent):
        """
        :param agent: agent that regulates the vehicle
        :param vehicle: actor to apply to local planner logic onto
        """
        self._vehicle = agent.vehicle
        self._sampling_resolution = agent._sampling_resolution
        self._dt = 1.0/self.FPS
        
        vehicle_physics_control = self._vehicle.get_physics_control()
        weels = vehicle_physics_control.wheels
        # self._front_centre = (weels[0].position + weels[1].position) / 100 / 2
        self._axis_length = compute_distance(weels[0].position, weels[2].position) / 100 # cm -> m
        self._max_steer_angle = weels[0].max_steer_angle / 2 # degree 假定内外角一致
        
        self._map = agent.vehicle.get_world().get_map()
        self._world = agent.vehicle.get_world()

        self._target_speed = None
        self.sampling_radius = None
        self._min_distance = None

        self._current_waypoint = None
        self._current_transform = None

        self.target_road_option = None
        self._next_waypoints = None
        self.target_waypoint = None
        self.local_transform = None # add 
        self._vehicle_controller = None
        self._global_plan = None
        self._pid_controller = None
        
        self.waypoints_queue = deque(maxlen=20000)  # queue with tuples of (waypoint, RoadOption) 10 km
        self.smooth_frenetpoints = deque(maxlen=20000) # queue with waypoint.transform
        
        self._buffer_size = math.ceil(50 / agent._sampling_resolution) # 0.5m res 50m/0.5 向上取整
        self._waypoint_buffer = deque(maxlen = self._buffer_size)
        self._transform_buffer = deque(maxlen = self._buffer_size)
        
        self._init_controller()  # initializing controller
        self.next_speed = -1

    def get_vehicle_frontcenter(self):
        """
        获取车辆前轴中心位置
        @param NONE
        @return location:Carla 3D vector
        """
        current_transform = self._vehicle.get_transform()
        location = current_transform.location
        yaw = current_transform.rotation.yaw * math.pi / 180
        location.x += math.cos(yaw) * self._axis_length / 2
        location.y += math.sin(yaw) * self._axis_length / 2
        return location

    def reset_vehicle(self):
        """Reset the ego-vehicle"""
        self._vehicle = None
        print("Resetting ego-vehicle!")

    def _init_controller(self):
        """
        Controller initialization.

        dt -- time difference between physics control in seconds.
        This is can be fixed from server side
        using the arguments -benchmark -fps=F, since dt = 1/F

        target_speed -- desired cruise speed in km/h

        min_distance -- minimum distance to remove waypoint from queue

        lateral_dict -- dictionary of arguments to setup the lateral PID controller
                            {'K_P':, 'K_D':, 'K_I':, 'dt'}

        longitudinal_dict -- dictionary of arguments to setup the longitudinal PID controller
                            {'K_P':, 'K_D':, 'K_I':, 'dt'}
        """
        # Default parameters
        self.args_lat_hw_dict = {
            'K_P': 0.75,
            'K_D': 0.02,
            'K_I': 0.4,
            'dt': 1.0 / self.FPS}
        self.args_lat_city_dict = {
            'K_P': 0.58,
            'K_D': 0.02,
            'K_I': 0.5,
            'dt': 1.0 / self.FPS}
        self.args_long_hw_dict = {
            'K_P': 0.37,
            'K_D': 0.024,
            'K_I': 0.032,
            'dt': 1.0 / self.FPS}
        self.args_long_city_dict = {
            'K_P': 0.15,
            'K_D': 0.05,
            'K_I': 0.07,
            'dt': 1.0 / self.FPS}

        self._current_waypoint = self._map.get_waypoint(self._vehicle.get_location())

        self._global_plan = False

        self._target_speed = self._vehicle.get_speed_limit()

        self._min_distance = self._sampling_resolution * 2

    def set_speed(self, speed):
        """
        Request new target speed.

            :param speed: new target speed in km/h
        """
        self._target_speed = speed

    def set_global_plan(self, current_plan):
        """
        将carla map的全局路径转换为pnc需要的全局路径：
        (transform, s, kappa, dappa)，保持 smooth_frenetpoints的序号与 current_plan一致
        """
        total_number = len(current_plan)
        change_lane_number = []
        for i in range(1, total_number):
            current_wp, _ = current_plan[i]
            previous_wp, _ = current_plan[i-1]
            # 两点距离，两点间角度变化
            p_c_dist, p_c_angle = compute_magnitude_angle(current_wp.transform.location, 
                                                                                                                      previous_wp.transform.location, 
                                                                                                                      previous_wp.transform.rotation.yaw) 
            diff_yaw = current_wp.transform.rotation.yaw - previous_wp.transform.rotation.yaw # 自有朝向变化

            delta = diff_yaw - p_c_angle # 
            delta = math.fabs(clip_angle(delta, -180, 180))
            # 转向变化过度：两点之间的夹角与两点朝向角变化的差值过大
            if delta > math.degrees(2 * self._sampling_resolution / 6.0) and p_c_dist > self._sampling_resolution / 2.0: # 最小转弯半径6.0米
                change_lane_number.append(i)

        # 前后选择10米重新计算location和rotation，10米是拍脑袋,假定的是换一个车道4米，
        mod_len = 10 
        jump_index = int(mod_len / self._sampling_resolution)
        calc_lane_number = []
        # 计算需要变换的起始点和结束点
        for index in change_lane_number:
            start_index = index - jump_index if index-jump_index>=0 else 0
            end_index = index+jump_index if index+jump_index<total_number else total_number - 1
            if calc_lane_number:
                if start_index < calc_lane_number[-1][1]:
                    start_index = calc_lane_number[-1][1]
                calc_lane_number.append((start_index, end_index))
            else:
                calc_lane_number.append((start_index, end_index)) 
        print(calc_lane_number)
        # 先填充        
        for i, elem in enumerate(current_plan):
            self.waypoints_queue.append(elem)
            self.smooth_frenetpoints.append([elem[0].transform, elem[0].s, 0.0, 0.0])
        # 计算bezier曲线
        for (si, ei) in calc_lane_number:
            # bezier
            loc0 = self.smooth_frenetpoints[si][0].location
            yaw0 = self.smooth_frenetpoints[si][0].rotation.yaw
            loc1 = self.smooth_frenetpoints[ei][0].location
            yaw1 = self.smooth_frenetpoints[ei][0].rotation.yaw
            number = ei - si + 1
            wave = GetWaveBezier(loc0.x, loc0.y, yaw0, loc1.x, loc1.y, yaw1, number)
            for i, (x,y) in enumerate(wave): # wave是bezier曲线
                self.smooth_frenetpoints[si+i][0].location.x = float(x) #loc_xy[i][0]
                self.smooth_frenetpoints[si+i][0].location.y = float(y) #loc_xy[i][0]
                if i > 0:
                    # 第一点朝向不变，第二点是1-2点线段的朝向
                    cur_wp = self.smooth_frenetpoints[si+i][0]
                    pre_wp = self.smooth_frenetpoints[si+i-1][0]
                    # dist, angle = compute_magnitude_angle(cur_wp.location, pre_wp.location, pre_wp.rotation.yaw)
                    dist, angle = compute_magnitude_angle(cur_wp.location, pre_wp.location, 0.0)
                    self.smooth_frenetpoints[si+i][0].rotation.yaw = -angle
                    self.smooth_frenetpoints[si+i][1] = self.smooth_frenetpoints[si+i-1][1] + dist
        # 更新s、kappa(曲率)、dappa(曲率变化率)等
        for i, _ in enumerate(self.smooth_frenetpoints):
            if i == 0:
                self.smooth_frenetpoints[i][1] = 0 # s
                self.smooth_frenetpoints[i][2] = 0 # 曲率
                self.smooth_frenetpoints[i][3] = 0 # 曲率变化率
            else: 
                cur_wp = self.smooth_frenetpoints[i][0]
                pre_wp = self.smooth_frenetpoints[i-1][0]
                dist, _ = compute_magnitude_angle(cur_wp.location, pre_wp.location, pre_wp.rotation.yaw)
                kappa = math.radians(cur_wp.rotation.yaw - pre_wp.rotation.yaw) / dist
                dappa = (kappa - self.smooth_frenetpoints[i-1][2]) / dist
                self.smooth_frenetpoints[i][1] = self.smooth_frenetpoints[i-1][1] + dist # s
                self.smooth_frenetpoints[i][2] = kappa # 曲率
                self.smooth_frenetpoints[i][3] = dappa # 曲率变化率

        # for debug
        _tfs = []       
        for (tf, s, kappa, dappa) in self.smooth_frenetpoints:
            _tfs.append(tf)
        draw_transforms(self._world, _tfs, time=-1)
        self._global_plan = True
        # for debug
        for i, (wp, _) in enumerate(self.waypoints_queue):
            print(wp.transform.location, '|', self.smooth_frenetpoints[i][0].location)
            print()

    def get_incoming_waypoint_and_direction(self, steps=3):
        """
        Returns direction and waypoint at a distance ahead defined by the user.
            :param steps: number of steps to get the incoming waypoint.
        """
        if len(self.waypoints_queue) > steps:
            return self.waypoints_queue[steps]

        else:
            try:
                wpt, direction = self.waypoints_queue[-1]
                return wpt, direction
            except IndexError as i:
                print(i)
                return None, RoadOption.VOID
        return None, RoadOption.VOID

    def run_step(self, target_speed=None, debug=True):
        """
        https://blog.csdn.net/u013468614/article/details/103518011
        δ(k)=θe(k)+arctan(λe(k)/v(k))，
        其中，θe(k)为k时刻的航向角偏差，计算方法：目标点方向 - 按现在控制参数车辆到达点的方向
        e(k)为根据上图几何学计算的横向跟踪误差，计算方法：按照当前控制参数车辆到达点 与 目标点 之间的距离
        λ为需要调节的参数，v为无人车当前速度，δ(k)为前轮偏角
        需要注意的是：因为方向盘转角速度是有限制的，因此δ(k)的变化率不能太大，却δ(k)是有最大值的

        stanley的思想是：当我以delta(k)角运动时，我可以完全补偿完所有的横向误差（假定到达点和当前点的朝向一致，即在一条直线上）

        带前馈的stanly模型：
        1.计算当前位置上的横向误差
        2.根据stanley算法计算前轮转角 delta0
        3.计算dt能到达的前馈点
        4.计算前馈点朝向与当前点朝向的差 delta1
        """
        control = carla.VehicleControl(steer=0, throttle=0) # init

        # 如果到达终点，则停止
        if len(self.waypoints_queue) == 0:
            control = carla.VehicleControl()
            control.steer = 0.0
            control.throttle = 0.0
            control.brake = 1.0
            control.hand_brake = False
            control.manual_gear_shift = False
            return control

        # 如果局部路径（_waypoint_buffer）为空，填充路径
        if not self._waypoint_buffer: # 空
            for i in range(self._buffer_size):
                if self.waypoints_queue:
                    self._waypoint_buffer.append(
                        self.waypoints_queue.popleft())
                else:
                    break            

        # 如果
        if not self._transform_buffer: # 空
            for i in range(self._buffer_size):
                if self.smooth_frenetpoints:
                    self._transform_buffer.append(
                        self.smooth_frenetpoints.popleft())
                else:
                    break            

        # 以下为控制逻辑
        # 计算前轮中心，以便执行stanley控制
        front_location = self.get_vehicle_frontcenter() # 前轮中心
        current_index = get_nearest_waypoint_index(self._waypoint_buffer, front_location) #
        current_index = get_nearest_location_index(self._transform_buffer, front_location)
        self._current_waypoint, _ = self._waypoint_buffer[current_index]
        self._current_transform, _, _, _ = self._transform_buffer[current_index]

        # for debug
        draw_waypoints(self._vehicle.get_world(), [self._current_waypoint], z=2.5) # 
        car_transform = self._vehicle.get_transform()
        car_transform.location = front_location

        # 获取本车当前状态
        speed = get_speed(self._vehicle) # 速度
        acc =get_acceleration(self._vehicle) # 加速度
        current_steer = self._vehicle.get_control().steer # 当前转角比例
        draw_speed_with_tf(self._vehicle.get_world(), car_transform, speed=speed, color=carla.Color(0,255,0), z=2.5) #

        # 计算当前状态下可能到达的target点，后期stanley中可以考虑增加预瞄点
        dist = speed*self._dt + acc / 2.0 * self._dt * self._dt + self._axis_length / 2 # 增加了半个轴长的预瞄
        ind = math.ceil(dist / self._sampling_resolution) # 向上取整
        if ind > len(self._waypoint_buffer):
            ind = -1
        self.target_waypoint, self.target_road_option = self._waypoint_buffer[ind]
        self.target_transform, _, _, _ = self._transform_buffer[ind]

        # 获取本车位置
        yaw = self._vehicle.get_transform().rotation.yaw * math.pi / 180.0
        
        # 初始化stanley横向控制器
        stanley_controller = StanleyController(front_location, yaw=yaw, v=speed, dt = self._dt, k=0.3)
        # 计算控制角度
        next_steer, lat_dist_error = stanley_controller.run_step(self._current_transform, self._max_steer_angle)
        # 考虑车辆角度变化的快慢能力，限制转角速度
        next_steer = np.clip(next_steer, -1, 1)
        d_steer = 0.05
        if next_steer > current_steer + d_steer:
            next_steer = current_steer + d_steer
        elif next_steer < current_steer - d_steer:
            next_steer = current_steer - d_steer

        # 准备进行纵向速度控制
        # pid参数选择
        if target_speed > 50:
            args_lat = self.args_lat_hw_dict
            args_long = self.args_long_hw_dict
        else:
            args_lat = self.args_lat_city_dict
            args_long = self.args_long_city_dict

        self._pid_controller = VehiclePIDController(self._vehicle,
                                                    args_lateral=args_lat,
                                                    args_longitudinal=args_long) # 横向pid参数，纵向pid参数

        # 计算3s内的道路曲率，从而限定预期的目标
        pre_aim = max(get_speed(self._vehicle) * 3.0, 9.0) # 预瞄位置：最小转弯半径6.0，按照6/0.7计算
        pre_indexs = math.ceil(pre_aim / self._sampling_resolution)
        max_curv = 0.0
        for i in range(1, min(len(self._waypoint_buffer), pre_indexs)):
            wp0, _ = self._waypoint_buffer[i-1]
            wp1, _ = self._waypoint_buffer[i]
            curv = math.fabs(get_curvature(wp0, wp1)) # 不考虑正负
            if curv > max_curv:
                max_curv = curv
        # 根据转角、横向误差、最大曲率来确定目标速度
        if math.fabs(next_steer) > 0.7 or lat_dist_error > 3.0 or max_curv > 1.0/8.0:
            self._target_speed = 3.0
        elif math.fabs(next_steer) > 0.4 or lat_dist_error > 1.5 or max_curv > 1.0/15.0:
            self._target_speed = 5.0
        elif math.fabs(next_steer) > 0.2 or lat_dist_error >  0.5 or max_curv > 1.0/25.0:
            self._target_speed = 6.0
        else:
            self._target_speed = 7.0
        control = self._pid_controller.run_step(self._target_speed, self.target_transform)

        # for debug 显示角度控制值
        control.steer = next_steer
        car_transform.rotation.yaw += next_steer*self._max_steer_angle
        draw_speed_with_tf(self._vehicle.get_world(), car_transform, speed=speed, color=carla.Color(255,255,0), z=2.5) #

        # Purge the queue of obsolete waypoints
        if current_index > 0:
            for i in range(current_index):
                self._waypoint_buffer.popleft()

                # 增加buffer taosheng 6.11
                if self.waypoints_queue:
                    self._waypoint_buffer.append(
                        self.waypoints_queue.popleft())

        if current_index > 0:
            for i in range(current_index):
                self._transform_buffer.popleft()

                # 增加buffer taosheng 6.11
                if self.smooth_frenetpoints:
                    self._transform_buffer.append(
                        self.smooth_frenetpoints.popleft())
        
        # for debug 显示局部规划路径
        if debug:
            tfs = list(zip(*self._transform_buffer))[0]
            if tfs:
                draw_transforms(self._world, tfs, time=0.1, thickness=0.5, color=carla.Color(0,0,255))
        return control


    def average_tracepoints(self, trace_points):
        _, new_ro = trace_points[0]
        new_wp = carla.Waypoint(carla.transform(carla.Location(0,0,0), carla.Rotation(0,0,0)))
        new_wp.transform.location.x = 0
        new_wp.transform.location.y = 0
        new_wp.transform.rotation.yaw = 0
        print('ori=', new_wp.transform.location.x,
            new_wp.transform.location.y,
            new_wp.transform.rotation.yaw)

        for (wp, _) in trace_points:
            new_wp.transform.location.x += wp.transform.location.x
            new_wp.transform.location.y += wp.transform.location.y
            new_wp.transform.rotation.yaw += wp.transform.rotation.yaw
        print('new=', new_wp.transform.location.x,
            new_wp.transform.location.y,
            new_wp.transform.rotation.yaw)
        new_wp.transform.location.x /= len(trace_points)
        new_wp.transform.location.y /= len(trace_points)
        new_wp.transform.rotation.yaw /= len(trace_points)
        print('div=', new_wp.transform.location.x,
            new_wp.transform.location.y,
            new_wp.transform.rotation.yaw)
        print('', len(trace_points))
        return (new_wp, new_ro)

    def smooth_trace(self, rough_route_trace, method='wm'):
        """
        平滑曲线 TODO
        """
        smooth_route = []
        if method=='wm':
            # 滑动窗口平均滤波，因为车道宽度是4米，考虑30度转角
            # 因此，滤波长度为8米
            # 滤波窗口点数
            wm = int(8.0 / self._sampling_resolution)
            print('wm=', wm)
            for i in range(len(rough_route_trace)):
                if i < wm:
                    new_tp = self.average_tracepoints(rough_route_trace[0:i+1])
                else:
                    new_tp = self.average_tracepoints(rough_route_trace[i-wm:i+1])
                smooth_route.append(new_tp)
        else:
            pass
        return smooth_route 
